"""
Flask后端主应用
AI智能助手 - 集成智谱清言所有API功能
"""
from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
import os
from werkzeug.utils import secure_filename
import json
from database import db
import config
from api_client import ZhipuAPI

app = Flask(__name__, template_folder='templates')
CORS(app)  # 允许跨域请求

# 配置Jinja2使用 [[ ]] 作为插值分隔符，避免与Vue.js的 {{ }} 冲突
app.jinja_env.variable_start_string = '[['
app.jinja_env.variable_end_string = ']]'

# 配置
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'pdf', 'docx', 'pptx', 'txt', 'md'}
MAX_FILE_SIZE = 16 * 1024 * 1024  # 16MB

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs('generated', exist_ok=True)
os.makedirs('audio', exist_ok=True)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = MAX_FILE_SIZE

# 初始化API客户端
api_client = ZhipuAPI(api_key=config.API_KEY, base_url=config.BASE_URL)


def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# ============ 前端页面 ============
@app.route('/')
def index():
    """前端首页"""
    return render_template('index.html')


# ============ 会话相关接口 ============
@app.route('/api/conversations', methods=['GET'])
def get_conversations():
    """获取所有会话列表"""
    try:
        conversations = db.get_conversations()
        return jsonify({'success': True, 'data': conversations})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/conversations', methods=['POST'])
def create_conversation():
    """创建新会话"""
    try:
        data = request.get_json()
        title = data.get('title', '新对话')
        conversation_id = db.create_conversation(title)
        return jsonify({'success': True, 'data': {'id': conversation_id, 'title': title}})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/conversations/<int:conversation_id>', methods=['GET'])
def get_conversation(conversation_id):
    """获取单个会话详情"""
    try:
        conversation = db.get_conversation(conversation_id)
        if not conversation:
            return jsonify({'success': False, 'error': '会话不存在'}), 404
        messages = db.get_messages(conversation_id)
        conversation['messages'] = messages
        return jsonify({'success': True, 'data': conversation})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/conversations/<int:conversation_id>', methods=['DELETE'])
def delete_conversation(conversation_id):
    """删除会话"""
    try:
        conversation = db.get_conversation(conversation_id)
        if not conversation:
            return jsonify({'success': False, 'error': '会话不存在'}), 404
        db.delete_conversation(conversation_id)
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/conversations/<int:conversation_id>/messages', methods=['DELETE'])
def clear_messages(conversation_id):
    """清除会话的所有消息（保留会话）"""
    try:
        conversation = db.get_conversation(conversation_id)
        if not conversation:
            # 如果会话不存在，返回错误（默认会话应该始终存在）
            return jsonify({'success': False, 'error': '会话不存在'}), 404
        # 清除消息
        db.clear_messages(conversation_id)
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


# ============ 对话相关接口 ============
@app.route('/api/chat', methods=['POST'])
def chat():
    """对话补全（流式输出）"""
    try:
        data = request.get_json()
        conversation_id = data.get('conversation_id', 1)  # 默认使用会话ID=1
        message = data.get('message')

        if not message:
            return jsonify({'success': False, 'error': '消息内容不能为空'}), 400

        # 确保会话存在（默认会话ID=1应该始终存在）
        conversation = db.get_conversation(conversation_id)
        if not conversation:
            # 如果会话不存在，返回错误（理论上不应该发生，因为默认会话会在数据库初始化时创建）
            return jsonify({'success': False, 'error': '会话不存在'}), 404

        # 获取历史消息
        messages_history = []
        db_messages = db.get_messages(conversation_id)
        for msg in db_messages:
            messages_history.append({
                'role': msg['role'],
                'content': msg['content']
            })

        # 添加用户消息
        messages_history.append({'role': 'user', 'content': message})
        db.add_message(conversation_id, 'user', message)

        # 获取模型和参数设置
        model = data.get('model', 'glm-4.5-flash')
        temperature = data.get('temperature', 1)
        max_tokens = data.get('max_tokens', 65536)

        # 流式输出
        from flask import Response, stream_with_context
        def generate():
            full_content = ""
            for chunk in api_client.chat_completion(messages_history, model=model, temperature=temperature, max_tokens=max_tokens):
                if not chunk.get('success'):
                    yield f"data: {json.dumps({'error': chunk.get('error', '未知错误')})}\n\n"
                    break
                
                data = chunk.get('data', {})
                if 'choices' in data and len(data['choices']) > 0:
                    delta = data['choices'][0].get('delta', {})
                    content = delta.get('content', '')
                    if content:
                        full_content += content
                        yield f"data: {json.dumps({'content': content, 'done': False})}\n\n"
                    
                    finish_reason = data['choices'][0].get('finish_reason')
                    if finish_reason == 'stop':
                        # 保存完整消息到数据库
                        if conversation_id and full_content:
                            db.add_message(conversation_id, 'assistant', full_content)
                        yield f"data: {json.dumps({'content': '', 'done': True, 'full_content': full_content})}\n\n"
                        break
            
            yield "data: [DONE]\n\n"
        
        return Response(stream_with_context(generate()), mimetype='text/event-stream')

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


# ============ 文件上传和解析接口 ============
# 已移除：文件上传与解析功能


# ============ 图像生成接口 ============
@app.route('/api/images/generate', methods=['POST'])
def generate_image():
    """生成图像"""
    try:
        data = request.get_json()
        prompt = data.get('prompt')
        conversation_id = data.get('conversation_id')

        if not prompt:
            return jsonify({'success': False, 'error': '提示词不能为空'}), 400

        # 获取模型和参数设置
        model = data.get('model', 'cogview-3-flash')
        size = data.get('size', '1024x1024')

        # 生成图像，自动保存到generated目录
        response = api_client.generate_image(prompt, model=model, size=size, save_to_file=True)

        if response.get('success'):
            image_url = response['data']['data'][0]['url']
            # 如果已保存到本地，使用本地路径
            if 'file_path' in response:
                image_url = f"/api/files/generated/{response['filename']}"
            
            gen_id = db.add_generation('image', prompt, conversation_id)
            db.update_generation(gen_id, result_url=image_url, status='completed')
            
            return jsonify({
                'success': True,
                'data': {
                    'image_url': image_url,
                    'generation_id': gen_id,
                    'local_path': response.get('file_path')
                }
            })
        else:
            return jsonify({'success': False, 'error': response.get('error')}), 500

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


# ============ 视频生成接口 ============
@app.route('/api/videos/generate', methods=['POST'])
def generate_video():
    """生成视频（异步）"""
    try:
        data = request.get_json()
        prompt = data.get('prompt')
        conversation_id = data.get('conversation_id')

        if not prompt:
            return jsonify({'success': False, 'error': '提示词不能为空'}), 400

        # 获取模型和参数设置
        model = data.get('model', 'cogvideox-flash')
        quality = data.get('quality', 'quality')
        with_audio = data.get('with_audio', True)
        size = data.get('size', '1920x1080')
        fps = data.get('fps', 30)

        # 生成视频（异步）
        response = api_client.generate_video(prompt, model=model, quality=quality, with_audio=with_audio, size=size, fps=fps)

        if response.get('success'):
            task_id = response['data'].get('id')
            gen_id = db.add_generation('video', prompt, conversation_id)
            if task_id:
                db.add_task(task_id, 'video_generation', conversation_id)
            
            return jsonify({
                'success': True,
                'data': {
                    'task_id': task_id,
                    'generation_id': gen_id,
                    'status': 'processing'
                }
            })
        else:
            return jsonify({'success': False, 'error': response.get('error')}), 500

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/tasks/<task_id>', methods=['GET'])
def get_task_status(task_id):
    """查询异步任务状态"""
    try:
        # 先查数据库
        task_info = db.get_task(task_id)
        if not task_info:
            return jsonify({'success': False, 'error': '任务不存在'}), 404

        # 如果任务还在处理中，查询API
        if task_info['status'] in ['pending', 'processing']:
            response = api_client.get_async_result(task_id)
            if response.get('success'):
                result_data = response['data']
                
                # 检查是否有视频结果（如果有video_result字段，说明任务完成）
                video_result = result_data.get('video_result', [])
                if video_result and len(video_result) > 0:
                    status = 'completed'
                    # 更新生成记录
                    if task_info['task_type'] == 'video_generation':
                        video_url = video_result[0].get('url')
                        cover_url = video_result[0].get('cover_image_url')
                        
                        # 自动下载视频到本地（类似于图片生成）
                        if video_url:
                            try:
                                download_response = api_client.download_video(video_url, save_to_file=True)
                                if download_response.get('success') and 'file_path' in download_response:
                                    # 使用本地路径
                                    video_url = f"/api/files/generated/{download_response['filename']}"
                                    result_data['local_path'] = download_response['file_path']
                                    result_data['filename'] = download_response['filename']
                            except Exception as e:
                                # 下载失败不影响主要结果，仍然使用远程URL
                                print(f"视频下载失败: {e}")
                                result_data['download_error'] = str(e)
                        
                        result_data['video_url'] = video_url
                        result_data['cover_image_url'] = cover_url
                        
                        # 更新数据库记录
                        db.update_task_status(task_id, 'completed', json.dumps(result_data))
                        
                        # 更新生成记录的状态和结果URL
                        conversation_id = task_info.get('conversation_id')
                        if conversation_id:
                            # 查找对应的生成记录
                            conn = db.get_connection()
                            cursor = conn.cursor()
                            cursor.execute(
                                "SELECT id FROM generations WHERE conversation_id = ? AND generation_type = 'video' ORDER BY id DESC LIMIT 1",
                                (conversation_id,)
                            )
                            gen_row = cursor.fetchone()
                            if gen_row:
                                gen_id = gen_row['id']
                                local_path = result_data.get('local_path')
                                db.update_generation(gen_id, result_url=video_url, local_path=local_path, status='completed')
                            conn.close()
                    else:
                        db.update_task_status(task_id, 'completed', json.dumps(result_data))
                elif 'error' in result_data:
                    status = 'failed'
                    db.update_task_status(task_id, 'failed', error_message=result_data.get('error'))
                else:
                    status = 'processing'
                
                return jsonify({
                    'success': True,
                    'data': {
                        'task_id': task_id,
                        'status': status,
                        'result': result_data
                    }
                })
            else:
                return jsonify({'success': False, 'error': response.get('error')}), 500

        # 任务已完成或失败，直接返回数据库中的结果
        result_data = json.loads(task_info['result']) if task_info['result'] else None
        
        # 如果任务已完成且是视频生成任务，检查是否有本地路径
        if task_info['status'] == 'completed' and task_info['task_type'] == 'video_generation' and result_data:
            # 从数据库的 generations 表中获取本地路径信息
            conversation_id = task_info.get('conversation_id')
            if conversation_id:
                conn = db.get_connection()
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT result_url, local_path FROM generations WHERE conversation_id = ? AND generation_type = 'video' ORDER BY id DESC LIMIT 1",
                    (conversation_id,)
                )
                gen_row = cursor.fetchone()
                if gen_row:
                    # 如果数据库中有本地路径，使用本地路径
                    if gen_row['result_url'] and gen_row['result_url'].startswith('/api/files/generated/'):
                        result_data['video_url'] = gen_row['result_url']
                        if gen_row['local_path']:
                            result_data['local_path'] = gen_row['local_path']
                            result_data['filename'] = os.path.basename(gen_row['local_path'])
                conn.close()
        
        return jsonify({
            'success': True,
            'data': {
                'task_id': task_id,
                'status': task_info['status'],
                'result': result_data
            }
        })

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/audio/speech', methods=['POST'])
def text_to_speech():
    """文本转语音"""
    try:
        data = request.get_json()
        text = data.get('text')
        conversation_id = data.get('conversation_id')

        if not text:
            return jsonify({'success': False, 'error': '文本内容不能为空'}), 400

        # 获取模型和参数设置
        model = data.get('model', 'cogtts')
        voice = data.get('voice', 'tongtong')
        response_format = data.get('response_format', 'wav')

        # 调用API，自动保存到audio目录
        response = api_client.text_to_speech(text, model=model, voice=voice, response_format=response_format, save_to_file=True)

        if response.get('success'):
            # 获取保存的文件路径
            filename = response.get('filename')
            if filename:
                audio_url = f'/api/audio/play/{filename}'
                db.add_audio_record('speech', text_content=text, 
                                  audio_path=response.get('file_path'), conversation_id=conversation_id)
                
                return jsonify({
                    'success': True,
                    'data': {'audio_url': audio_url, 'filename': filename}
                })
            else:
                # 如果没有保存文件，返回音频数据
                audio_data = response['data']
                audio_filename = f"speech_{conversation_id or 'temp'}_{int(os.urandom(4).hex(), 16)}.wav"
                audio_path = os.path.join('audio', audio_filename)
                with open(audio_path, 'wb') as f:
                    f.write(audio_data)
                
                db.add_audio_record('speech', text_content=text, 
                                  audio_path=audio_path, conversation_id=conversation_id)
                
                return jsonify({
                    'success': True,
                    'data': {'audio_url': f'/api/audio/play/{audio_filename}'}
                })
        else:
            return jsonify({'success': False, 'error': response.get('error')}), 500

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/audio/play/<filename>', methods=['GET'])
def play_audio(filename):
    """播放音频文件"""
    try:
        audio_path = os.path.join('audio', filename)
        if os.path.exists(audio_path):
            return send_file(audio_path, mimetype='audio/mpeg', as_attachment=False)
        else:
            return jsonify({'success': False, 'error': '音频文件不存在'}), 404
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/files/generated/<filename>', methods=['GET'])
def get_generated_file(filename):
    """获取生成的文件（图片、视频等）"""
    try:
        # 安全检查：只允许访问generated目录下的文件
        filename = secure_filename(filename)
        file_path = os.path.join('generated', filename)
        
        if not os.path.exists(file_path):
            return jsonify({'success': False, 'error': '文件不存在'}), 404
        
        # 根据文件扩展名确定MIME类型
        ext = filename.rsplit('.', 1)[1].lower() if '.' in filename else ''
        mimetype_map = {
            'png': 'image/png',
            'jpg': 'image/jpeg',
            'jpeg': 'image/jpeg',
            'gif': 'image/gif',
            'webp': 'image/webp',
            'mp4': 'video/mp4',
            'webm': 'video/webm',
            'mov': 'video/quicktime'
        }
        mimetype = mimetype_map.get(ext, 'application/octet-stream')
        
        return send_file(file_path, mimetype=mimetype, as_attachment=False)
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


# 已移除：网络搜索接口


# 已移除：文件解析结果查询接口


# ============ 设置相关接口 ============
@app.route('/api/settings', methods=['GET'])
def get_settings():
    """获取所有设置"""
    try:
        settings = db.get_all_settings()
        return jsonify({'success': True, 'data': settings})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/settings/<setting_key>', methods=['GET'])
def get_setting(setting_key):
    """获取单个设置"""
    try:
        setting = db.get_setting(setting_key)
        if setting:
            return jsonify({'success': True, 'data': setting['setting_value']})
        else:
            return jsonify({'success': False, 'error': '设置不存在'}), 404
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/settings/<setting_key>', methods=['POST', 'PUT'])
def save_setting(setting_key):
    """保存单个设置"""
    try:
        data = request.get_json()
        setting_value = data.get('value')
        
        if setting_value is None:
            return jsonify({'success': False, 'error': '设置值不能为空'}), 400
        
        success = db.save_setting(setting_key, setting_value)
        if success:
            return jsonify({'success': True})
        else:
            return jsonify({'success': False, 'error': '保存设置失败'}), 500
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/settings', methods=['POST', 'PUT'])
def save_all_settings():
    """批量保存设置"""
    try:
        data = request.get_json()
        settings = data.get('settings')
        
        if not settings or not isinstance(settings, dict):
            return jsonify({'success': False, 'error': '设置数据格式错误'}), 400
        
        success = db.save_all_settings(settings)
        if success:
            return jsonify({'success': True})
        else:
            return jsonify({'success': False, 'error': '保存设置失败'}), 500
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


if __name__ == '__main__':
    print("AI智能助手后端启动中...")
    print(f"数据库路径: {db.db_path}")
    app.run(debug=True, port=5000)

